﻿#include <QFileDialog>
#include <detect_thread.h>
#include <QMessageBox>
#include <QThread>
#include <opencv2/opencv.hpp>

#include "commons.h"
#include "inference.h"

bool DetectThreadBase::STOP = false;
int DetectThreadBase::ID = 0;
QThread* DetectThreadBase::cur_thread = nullptr;
DetectThreadBase::DetectThreadBase(MainWindow& mw) : main_window(&mw)
{
}


void DetectThreadBase::stop()
{
	DetectThreadBase::STOP = true;
}

void DetectThreadBase::run() {
	this->before_running();
	this->do_detect();
	this->after_running();
}


void DetectThreadBase::before_running()
{
	DetectThreadBase::ID = this->get_detect_id();
	DetectThreadBase::STOP = false;
	DetectThreadBase::cur_thread = this;
}

void DetectThreadBase::after_running()
{
	DetectThreadBase::ID = 0;                // 重置ID未0，表示没有线程在检测
	DetectThreadBase::STOP = true;		     // 停止检测
	DetectThreadBase::cur_thread = nullptr;  
}

void DetectThreadBase::clear()
{
}


const int ImageDetectThread::ID = 1;
ImageDetectThread::ImageDetectThread(MainWindow& mw, const QString& file) :DetectThreadBase(mw), file(file) {

}
int ImageDetectThread::get_detect_id()
{
	return ImageDetectThread::ID;
}
void ImageDetectThread::do_detect() {
	auto radio_button = this->main_window->ui.islabel;
	if (this->main_window->inference == nullptr) {
		QMessageBox::warning(this->main_window, "Error", "未检测到模型");
		return;
	}
	cv::Mat frame = cv::imread(this->file.toStdString());
	this->main_window->inference->doPredict(frame, radio_button->isChecked());
	QPixmap pixmap = commons::cvMatToQPixmap(frame);
	this->main_window->show_img(pixmap);
}


const int VideoDetectThread::ID = 2;
VideoDetectThread::VideoDetectThread(MainWindow& mw, const QString& file) :DetectThreadBase(mw), file(file) {

}
int VideoDetectThread::get_detect_id()
{
	return VideoDetectThread::ID;
}
void VideoDetectThread::do_detect() {
	emit this->main_window->start_loading_movie_signal();
	cv::VideoCapture cap;
	cap.open(this->file.toStdString());
	emit this->main_window->stop_loading_movie_signal();


	int idx = 0;
	auto radio_button = this->main_window->ui.islabel;
	while (!DetectThreadBase::STOP)
	{
		idx++;
		if (idx % 2 == 0) {
			idx = 0;
			continue;
		}  
		// 获取新的一帧;
		cv::Mat frame;
		cap >> frame;
		if (frame.empty())
			return;
		// 显示新的帧;
		this->main_window->inference->doPredict(frame, radio_button->isChecked());
		QPixmap map = commons::cvMatToQPixmap(frame);
		this->main_window->show_img(map);
	}
	cap.release();
}

const int CameraDetectThread::ID = 3;
int CameraDetectThread::get_detect_id()
{
	return CameraDetectThread::ID;
}


void CameraDetectThread::do_detect() {

	emit this->main_window->start_loading_movie_signal();
	cv::VideoCapture cap;
	cap.open(0);
	if (!cap.isOpened())
		return;
	int idx = 0;
	emit this->main_window->stop_loading_movie_signal();

	auto radio_button = this->main_window->ui.islabel;
	while (!DetectThreadBase::STOP)
	{
		idx++;
		if (idx % 2 == 0) {
			idx = 0;
			continue;
		}
		cv::Mat frame;
		cap >> frame;
		if (frame.empty())
			return;
		cv::flip(frame, frame, cv::RotateFlags::ROTATE_180);
		this->main_window->inference->doPredict(frame, radio_button->isChecked());
		QPixmap map = commons::cvMatToQPixmap(frame);
		this->main_window->show_img(map);
	}
	cap.release();
	this->main_window->ui.show_detect->clear();
}

WarmUpThread::WarmUpThread(MainWindow& mw, const QString& path, float iou_thres, float conf_thres, const cv::Size& input_shape)
	:main_window(&mw), path(path), iou_thres(iou_thres), conf_thres(conf_thres), input_shape(input_shape)
{
}

void WarmUpThread::run()
{
	auto& inference = this->main_window->inference;
	if (inference == nullptr)
	{
		inference.reset(new Inference());
	}
	emit this->main_window->start_loading_movie_signal();
	inference->loadOnnxNetwork(this->path.toStdString(), this->input_shape, true);
	inference->modelNMSThreshold = this->iou_thres;
	inference->modelScoreThreshold = this->conf_thres;
	cv::Mat zero = cv::Mat::zeros(this->input_shape, CV_8UC3);
	inference->doPredict(zero);
	emit this->main_window->stop_loading_movie_signal();
}
